import gym
import numpy as np
from gym import spaces
import torch
import random
import matplotlib.pyplot as plt
import numpy as np

class SoftwareUpgradeEnv(gym.Env):
    def __init__(self, num_components, time_windows, dependency_matrix, jump_matrix, device):
        super(SoftwareUpgradeEnv, self).__init__()

        self.num_components = num_components  # 组件数量
        self.time_windows = time_windows      # 每个时间窗的长度
        self.current_window = 0               # 当前时间窗索引
        self.remaining_time = time_windows[0] # 当前时间窗剩余时间
        self.dependency_matrix = dependency_matrix      # 依赖矩阵(三维)  -1 表示无依赖, -2 表示无该版本
        self.jump_matrix = jump_matrix                  # 跳跃矩阵       控制组件是否允许跳跃升级 (1 表示允许跳跃，0 表示不允许)
        self.max_versions = [len(dependency_matrix[i, 0]) for i in range(num_components)]  # 每个组件最大版本数
        self.device = device
        self.max_time_window = max(time_windows)
        self.total_time = sum(time_windows)
        self.gantt_chart = UpgradeGanttChart(num_components, time_windows)  #可视化类
        self.current_time = 0

        # 组件当前版本状态
        self.state = np.zeros(num_components, dtype=int)
        # **预计算所有合法状态**
        self.valid_states = self.generate_valid_states()
        self.state_to_index = {state: idx for idx, state in enumerate(self.valid_states)}


        # **升级时间矩阵**（组件不同版本升级所需时间）
        #upgrade_time[c, a, b]：组件 c 从版本 a 升级到 b 所需的时间。
        self.upgrade_time = np.random.randint(1, 5, size=(num_components, max_versions, max_versions))

        # **裁剪后的动作空间**
        self.valid_actions = self.get_all_valid_actions()
        self.valid_actions.append((0, 0, 0))  # 增加一个只消耗1时间单位的动作

        # **状态空间**
        # **状态空间** (状态索引 + 归一化剩余时间)
        self.observation_space = spaces.Box(
            low=np.array([0, 0.0]), 
            high=np.array([len(self.valid_states) - 1, 1.0]), 
            dtype=np.float32
        )

        # **动作空间**
        self.action_space = spaces.Discrete(len(self.valid_actions))

    
    def get_current_valid_action_indices(self):
        """ 获取当前状态下合法的动作索引 """
        valid_indices = []
        for idx, (comp, v1, v2) in enumerate(self.valid_actions):
            if self.state[comp] == v2:  # 已是目标版本，跳过
                continue
            if self.state[comp] != v1:  # 当前版本不匹配，跳过
                continue
            if not self.is_dependency_satisfied(comp, v2):  # 依赖未满足，跳过
                continue
            # 确保该动作不会超出剩余时间
            upgrade_time = self.upgrade_time[comp, v1, v2]
            if upgrade_time > self.remaining_time:
                continue  # 该动作不可执行

            valid_indices.append(idx)
        valid_indices.append(len(self.valid_actions) - 1)
        return valid_indices


    def get_all_valid_actions(self):
        """ 预计算所有合法的升级动作，满足依赖关系和跳跃约束 """
        valid_actions = []
        for comp in range(self.num_components):
            for v1 in range(self.max_versions[comp]):
                for v2 in range(self.max_versions[comp]):
                    if v1 < v2:
                        # 1️⃣ **检查该版本是否可用 (D[i, i, v] != -2)**
                        if self.dependency_matrix[comp, comp, v2] == -2:
                            continue
                        
                        # 2️⃣ **允许跳跃**
                        if not (self.jump_matrix[comp, v1, v2] == 1 or v2 == v1 + 1):
                            continue

                        valid_actions.append((comp, v1, v2))

        return valid_actions


    def is_valid_state(self, state):
        """检查状态是否满足所有依赖关系"""
        for c in range(self.num_components):
            version = state[c]
            for d in range(self.num_components):
                required_version = self.dependency_matrix[c, d, version]
                if required_version >= 0 and state[d] < required_version:
                    return False  # 依赖未满足
        return True

    def generate_valid_states(self):
        """预生成所有可能的合法状态"""
        all_states = np.array(np.meshgrid(*[range(v) for v in self.max_versions])).T.reshape(-1, self.num_components)
        valid_states = [tuple(state) for state in all_states if self.is_valid_state(state)]
        return valid_states

    def is_dependency_satisfied(self, component, target_version):
        """ 检查目标版本是否满足所有组件依赖 """
        for j in range(self.num_components):
            required_version = self.dependency_matrix[component, j, target_version]
            if required_version >= 0 and (self.state[j] < required_version or self.dependency_matrix[j, j, self.state[j]] == -2):
                return False  # 依赖组件未满足，或者依赖的组件版本无效
        return True
    

    def step(self, action_index, is_need_visualization = False):
        """ 执行动作，并更新环境 """
        self.last_action_index = action_index
        if action_index == len(self.valid_actions) - 1:  # 如果是“什么都不做”的动作
            if self.remaining_time <= 2:  # 如果时间窗快结束，给予正奖励
                reward = 2  
            else:
                reward = -2
            done = False
            self.remaining_time -= 1  # 消耗1个时间单位
            self.current_time += 1
            if self.remaining_time <= 0:
                self.current_window += 1
                if self.current_window < len(self.time_windows):
                    self.remaining_time = self.time_windows[self.current_window]
                else:
                    done = True  # 所有时间窗用完，任务结束
                    reward = -100
            return self.get_observation(), reward, done, {}


        component, current_version, target_version = self.valid_actions[action_index]

        # 检查当前组件是否已在目标版本
        if self.state[component] == target_version:
            return self.get_observation(), -1, False, {}
        
        # 检查当前组件是否不等于动作原始版本
        if self.state[component] != current_version:
            return self.get_observation(), -1, False, {}

        # 检查依赖是否满足
        if not self.is_dependency_satisfied(component, target_version):
            return self.get_observation(), -1, False, {}

        # 获取升级所需时间
        upgrade_time = self.upgrade_time[component, current_version, target_version]

        # 检查时间窗约束
        if upgrade_time > self.remaining_time:
            reward = -3  # 超出时间窗
            done = False
        else:
            # 执行升级
            self.state[component] = target_version
            self.remaining_time -= upgrade_time
            if self.remaining_time == 0:
                reward = 6  # 完美填满时间窗，给予更高奖励
            else:
                reward = 4  

            if is_need_visualization:
                component, from_version, to_version = self.valid_actions[action_index]
                self.gantt_chart.record_upgrade(component, from_version, to_version, self.current_time, upgrade_time)
                self.current_time += upgrade_time

            # 如果当前时间窗用完，切换到下一个时间窗
            if self.remaining_time <= 0:
                self.current_window += 1
                if self.current_window < len(self.time_windows):
                    self.remaining_time = self.time_windows[self.current_window]
                else:
                    done = True  # 所有时间窗用完，任务结束
                    reward = -100
                    return self.get_observation(), reward, done, {}
                
            # 如果所有组件升级到最高级，任务结束
            if np.all(self.state == (np.array(self.max_versions) - 1)):  
                reward = 100
                done = True  
                return self.get_observation(), reward, done, {}

            done = False



        
        return self.get_observation(), reward, done, {}

    def get_observation(self):
        """返回 (状态索引, 归一化剩余时间)"""
        state_idx = self.state_to_index[tuple(self.state)]
        time_norm = self.remaining_time / self.max_time_window
        return np.array([state_idx, time_norm], dtype=np.float32)
    

    def reset(self):
        self.current_window = 0
        self.current_time = 0
        self.remaining_time = self.time_windows[0]
        self.state = np.zeros(self.num_components, dtype=int)
        state_idx = self.state_to_index[tuple(self.state)]
        return self.get_observation()


    def render(self, done):
        print(f"State: {self.state}, Remaining Time: {self.remaining_time}, \
            Current Window: {self.current_window}, Last Action: {self.valid_actions[self.last_action_index]}")
        if done:
            self.gantt_chart.generate_gantt_chart()



class UpgradeGanttChart:
    def __init__(self, num_components, time_windows):
        self.upgrade_records = []  # 记录升级信息
        self.num_components = num_components  # 组件数
        self.time_windows = time_windows  # 记录时间窗
        self.window_boundaries = self.calculate_window_boundaries()  # 计算时间窗边界
        self.min_display_time = 0.2  # 确保短时间升级可见

    def calculate_window_boundaries(self):
        """ 计算每个时间窗的开始和结束时间 """
        boundaries = [0]
        for t in self.time_windows:
            boundaries.append(boundaries[-1] + t)
        return boundaries

    def record_upgrade(self, component, from_version, to_version, start_time, duration):
        """ 记录升级信息，并确保最小绘制时间 """
        duration = max(duration, self.min_display_time)  # 避免过短时间被忽略
        end_time = start_time + duration
        self.upgrade_records.append((component, from_version, to_version, start_time, end_time))

    def generate_gantt_chart(self):
        """ 绘制完整的甘特图，并标注时间窗分界 """
        if not self.upgrade_records:
            print("No upgrades recorded.")
            return
        
        fig, ax = plt.subplots(figsize=(12, max(6, self.num_components * 0.5)))  # 纵向更紧凑

        component_height = 0.5  # 调整每个组件的条形高度，使纵轴更紧凑

        # 绘制时间窗背景
        for i in range(len(self.window_boundaries) - 1):
            if i % 2 == 0:
                ax.axvspan(self.window_boundaries[i], self.window_boundaries[i + 1], color='gray', alpha=0.1)

        # 绘制升级条（**严格对齐，不偏移**）
        for record in self.upgrade_records:
            component, from_version, to_version, start_time, end_time = record

            ax.barh(component, end_time - start_time, left=start_time, height=component_height, 
                    color='skyblue', edgecolor='black')

            ax.text((start_time + end_time) / 2, component, f"v{from_version} → v{to_version}", 
                    ha='center', va='center', color='black', fontsize=10, fontweight='bold')

        # 绘制时间窗分隔线
        for boundary in self.window_boundaries[1:-1]:  # 不绘制起点和终点
            ax.axvline(boundary, color='red', linestyle='--', linewidth=1.5, label="Time Window Boundary")

        # 设置 X 轴和 Y 轴
        total_time = self.window_boundaries[-1]
        ax.set_xlim(0, total_time)
        ax.set_xlabel("Time")
        ax.set_yticks(range(self.num_components))
        ax.set_yticklabels([f"Component {c}" for c in range(self.num_components)])
        ax.set_title("Software Upgrade Gantt Chart (With Time Windows)")
        plt.grid(axis='x', linestyle='--', alpha=0.5)
        plt.legend()
        plt.show()


# 定义组件数量、最大版本数、时间窗长度
num_components = 5
max_versions = 10
time_windows = [10, 15, 12, 10]  # 每个时间窗的长度

# device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
device = torch.device("cpu")



# 定义依赖矩阵（随机生成，实际应用应自定义）
dependency_matrix = np.full((num_components, num_components, max_versions), -1)
for i in range(num_components):
    for v in range(max_versions):
        dependency_matrix[i, i, v] = 0  # 自己对自己无依赖

# 定义跳跃矩阵（随机生成，实际应用应自定义）
jump_matrix = np.random.choice([0, 1], size=(num_components, max_versions, max_versions))

# 创建环境
env = SoftwareUpgradeEnv(num_components, time_windows, dependency_matrix, jump_matrix, device)

# 测试运行


if __name__ == "__main__":
    obs = env.reset()
    done = False
    while not done:
        # action = env.action_space.sample()
        action = env.get_current_valid_action_indices()
        obs, reward, done, _ = env.step(random.choice(action), True)
        env.render(done)

        